{-# LANGUAGE FlexibleInstances, FlexibleContexts #-}

module Image ( 
    Image(..), ImageI, ImageD, 

    mkImage,

    modifyPixel, modifyPixels,
    readPixel,

--    imgToPngFile,
    imgToFits,
) where

-- import qualified          Graphics.GD as GD
import qualified Data.ByteString.Lazy as L

import               Fits
import      Control.Monad ( forM_ )
import           Data.Int ( Int16, Int64 )
import      Data.Array.IO ( MArray(..), IOUArray(..), readArray, writeArray, freeze )
import Data.Array.Unboxed ( UArray(..), elems )
import    Data.Binary.Put ( runPut )
import        Data.Binary ( Binary(..), put)

type Width   = Int
type Height  = Int

data Image a = Image { w   :: Width,
                       h   :: Height,
                       img :: IOUArray Int a}


type ImageI = Image Int64

type ImageD = Image Double

                      
mkImage :: (MArray IOUArray a IO) => Width -> Height -> a -> IO (Image a)
mkImage w h v =
    do
      arr <- newArray (1, w*h) v
      return Image { w = w, h = h, img = arr }


imgToByteString :: ImageD -> IO L.ByteString
imgToByteString (Image _ _ arr) = 
   do
      arr' <- freeze arr :: IO (UArray Int Double)

      return $ runPut $ mapM_ put (elems arr')

imgToByteStringWith :: (Binary a) => (Double -> a) -> ImageD -> IO L.ByteString
imgToByteStringWith expr (Image _ _ arr) = 
   do
      arr' <- freeze arr :: IO (UArray Int Double)

      let es = map expr (elems arr')

      return $ runPut $ mapM_ put es


modifyPixel :: (MArray IOUArray a IO) => Image a -> (a -> b -> a) -> ((Int,Int), b) -> IO ()
modifyPixel (Image w h arr) expr ((x,y), v) =
    if x<1 || x>w || y<1 || y>h then
      return ()
    else
      do 
        c <- readArray arr i
        writeArray arr i (expr c v)
    where
       i = (y - 1) * w + x


modifyPixels img expr ps = forM_ ps $ modifyPixel img expr      

readPixel :: (MArray IOUArray a IO) => Image a -> (Int, Int) -> IO a
readPixel (Image w h arr) (x,y) =
    do
      readArray arr i
    where
      i = (y - 1) * w + x

{-
getColor :: Int64 -> GD.Color
getColor i = GD.rgba c c c 80
    where
      c  = fromIntegral $ div i 16

imgToPngFile :: FilePath -> ImageI -> IO()
imgToPngFile file (Image w h arr) =
    do
      pic <- GD.newImage  (w, h)

      forM_ [(x,y) | x <-[1..w], y <- [1..h] ] $ drowPixel pic

      GD.savePngFile file pic

    where
      drowPixel pic (x,y) =
        do
          c <- readArray arr $ (y-1) * w + x
          GD.setPixel (x-1,y-1) (getColor c) pic
-}

imgToFits :: FilePath -> ImageD -> IO()
imgToFits file img =
    do
      hdu <- serializeImage img

      writeFits file hdu


instance FitsImage (Image Double) where
    serializeImage img@(Image w h _) =
        do
          dat <- imgToByteStringWith (\x ->  (round x :: Int16) - 2^15 ) img

          let head = [ Record (Keyword "SIMPLE") (VBool True)  "file conforms to FITS standard"
                     , Record (Keyword "BITPIX") (VInt 16)     "number of bits per data pixel"
                     , Record (Keyword "NAXIS")  (VInt 2)      "number of data axes"
                     , Record (Keyword "NAXIS1") (VInt w)      "length of data axis 1"
                     , Record (Keyword "NAXIS2") (VInt h)      "length of data axis 2"
                     , Record (Keyword "BZERO")  (VInt (2^15)) "physical value = BZERO + array value"
                     , END
                     ]
          return (HDUImage head dat)
